package drds.plus.datanode.select.equity_manager_select;

import drds.plus.datanode.configuration.DataSourceWrapper;
import drds.plus.datanode.executor.Executor;
import drds.plus.datanode.executor.ExecutorWrapper;
import drds.plus.datanode.select.AbstractSelector;
import drds.plus.datanode.select.DataSourceHolder;
import drds.plus.datasource.api.DataSource;
import lombok.extern.slf4j.Slf4j;

import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * <pre>
 * 按优先级选择的selector 每次选择只从优先级最高的一组DB中选择，若都不可用，才继续在下一个优先级的DB组中选择,优先级相同的DB还用随机选择
 *
 * 原始需求：TC要求在每个dbgroup中优先读备库，当备库不可用时，自动读主库
 * 扩展需求：一主多备，优先随机读备库。当备库都不可用时，才读主库
 *
 * 为了方便处理和接口一致，有如下要求：
 * 1. 目前只支持读分优先级组
 * 2. 一个权重推送的信息中，。。。
 * 3. 一个数据源只能在一个优先级组中？
 * </pre>
 */
@Slf4j
public class EquityManagerPrioritySelector extends AbstractSelector {

    /**
     * 按优先级顺序存放数据库组。元素0优先级最高。每个EquityDbManager元素代表具有相同优先级的一组数据库
     */
    private EquityManager[] equityManagers;

    public EquityManagerPrioritySelector(EquityManager[] equityManagers) {
        this.equityManagers = equityManagers;
        if (equityManagers == null || equityManagers.length == 0) {
            throw new IllegalArgumentException("EquityManager[] equityManagers is null or empty");
        }
    }

    public DataSourceWrapper select() {
        for (int i = 0; i < equityManagers.length; i++) {
            DataSourceWrapper dataSourceWrapper = equityManagers[i].select();
            if (dataSourceWrapper != null) {
                return dataSourceWrapper;
            }
        }
        return null;
    }

    public DataSourceWrapper get(String dataSourceId) {
        for (int i = 0; i < equityManagers.length; i++) {
            DataSourceWrapper dataSourceWrapper = equityManagers[i].get(dataSourceId);
            if (dataSourceWrapper != null) {
                return dataSourceWrapper;
            }
        }
        return null;
    }

    /**
     * 取每个级别的weightKey和总的weightKey的交集，挨个设置
     */
    public void setWeightConfigMap(Map<String, Integer> weightMap) {
        for (int i = 0; i < equityManagers.length; i++) {
            Map<String, Integer> oldWeightConfigMap = equityManagers[i].getRandomWeightSelect().getWeightConfigMap();
            Map<String, Integer> newWeightConfigMap = new HashMap<String, Integer>(oldWeightConfigMap.size());
            for (Map.Entry<String, Integer> entry : weightMap.entrySet()) {
                if (oldWeightConfigMap.containsKey(entry.getKey())) {
                    newWeightConfigMap.put(entry.getKey(), entry.getValue());
                }
            }
            equityManagers[i].setRandomWeightSelect(new RandomWeightSelect(newWeightConfigMap));
        }

    }

    /**
     * 基于EquityDbManager的tryExecute实现，对用户的tryer做一个包装，在wrapperTryer. onSQLException中
     * 检测到最后一个e是NoMoreDataSourceException时，不调原tryer的onSQLException, 转而重试其他优先级的
     */
    protected <T> T selectDataSourceHolderAndGetConnectionWrapperAndExecuteAndRetryIfFailed(Map<javax.sql.DataSource, SQLException> failedDataSources, Executor<T> executor, int times, Object... args) throws SQLException {
        final List<SQLException> historyExceptions = new ArrayList<SQLException>(0);
        Executor<T> wrapperTryer = new ExecutorWrapper<T>(executor, historyExceptions); // 移花接木

        for (int i = 0; i < equityManagers.length; i++) {
            try {
                return equityManagers[i].getConnectionWrapperAndExecuteAndRetryIfFailed(failedDataSources, wrapperTryer, times, args);
            } catch (GroupNotAvaliableException e) {
                log.warn("NoMoreDataSource for executor for priority datanode " + i);
            }
        }
        // 所有的优先级组都不可用，则抛出异常
        return executor.onSqlException(historyExceptions, exceptionSorter, args);
    }

    public void setSupportRetry(boolean isSupportRetry) {
        for (int i = 0; i < equityManagers.length; i++) {
            equityManagers[i].setSupportRetry(isSupportRetry);
        }
        this.isSupportRetry = isSupportRetry;
    }

    public void setReadable(boolean readable) {
        for (int i = 0; i < equityManagers.length; i++) {
            equityManagers[i].setReadable(readable);
        }

        this.readable = readable;
    }

    public Map<String, DataSource> getDataSourceIdToDataSourceHolderMap() {
        throw new UnsupportedOperationException("getDataSourceIdToDataSourceHolderMap()");
    }


    protected DataSourceHolder findDataSourceHolderByDataSourceIndex(String dataSourceIndex) {
        for (int i = 0; i < equityManagers.length; i++) {
            DataSourceHolder dataSourceHolder = equityManagers[i].findDataSourceHolderByDataSourceIndex(dataSourceIndex);
            if (dataSourceHolder != null) {
                return dataSourceHolder;
            }

        }
        return null;
    }

}
